import logging
import time
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, Optional, Set, Tuple
from uuid import uuid4

import redis.asyncio as redis
from jose import jwt

logger = logging.getLogger(__name__)


class SecurityUtils:
    # JWT 标准字段常量
    JWT_RESERVED_FIELDS: Set[str] = {"exp", "jti", "type", "iat", "nbf", "iss", "aud"}

    def __init__(
        self,
        secret_key: str,
        algorithm: str = "HS256",
        access_token_expire_minutes: int = 30,
        refresh_token_expire_minutes: int = 60 * 24 * 7,  # 7天
        redis_client: Optional[redis.Redis] = None,
        blacklist_prefix: str = "token_blacklist",
    ):
        """
        异步安全工具类（支持 Redis 黑名单）

        :param secret_key: JWT 签名密钥
        :param algorithm: JWT 算法（如 HS256）
        :param access_token_expire_minutes: Access Token 过期时间（分钟）
        :param refresh_token_expire_minutes: Refresh Token 过期时间（分钟）
        :param redis_client: 异步 Redis 客户端（若为 None，则黑名单功能禁用）
        :param blacklist_prefix: Redis key 前缀，避免冲突
        :raises ValueError: 当 secret_key 为空时抛出
        """
        if not secret_key:
            raise ValueError("secret_key is required and cannot be empty")

        self.secret_key = secret_key
        self.algorithm = algorithm
        self.access_token_expire_minutes = access_token_expire_minutes
        self.refresh_token_expire_minutes = refresh_token_expire_minutes
        self.redis = redis_client
        self.blacklist_prefix = blacklist_prefix

    def _create_token(self, data: Dict[str, Any], expires_delta: timedelta, token_type: str = "access") -> str:
        """
        创建 JWT token

        :param data: 要编码的数据
        :param expires_delta: 过期时间差
        :param token_type: token 类型（access 或 refresh）
        :return: 编码后的 JWT token
        """
        to_encode = data.copy()
        expire = datetime.now(timezone.utc) + expires_delta
        jti = str(uuid4())
        to_encode.update({"exp": expire, "jti": jti, "type": token_type})
        return jwt.encode(to_encode, self.secret_key, algorithm=self.algorithm)

    def create_access_token(self, data: Dict[str, Any], expires_delta: Optional[timedelta] = None) -> str:
        """
        创建访问 token

        :param data: 要编码的数据
        :param expires_delta: 过期时间差，如果为 None 则使用默认值
        :return: 访问 token
        """
        delta = expires_delta or timedelta(minutes=self.access_token_expire_minutes)
        return self._create_token(data, delta, "access")

    def create_refresh_token(self, data: Dict[str, Any], expires_delta: Optional[timedelta] = None) -> str:
        """
        创建刷新 token

        :param data: 要编码的数据
        :param expires_delta: 过期时间差，如果为 None 则使用默认值
        :return: 刷新 token
        """
        delta = expires_delta or timedelta(minutes=self.refresh_token_expire_minutes)
        return self._create_token(data, delta, "refresh")

    def create_tokens(self, data: Dict[str, Any]) -> Tuple[str, str]:
        """
        同时创建访问 token 和刷新 token

        :param data: 要编码的数据
        :return: (访问 token, 刷新 token) 的元组
        """
        access = self.create_access_token(data)
        refresh = self.create_refresh_token(data)
        return access, refresh

    async def _is_token_blacklisted(self, payload: Dict[str, Any]) -> bool:
        """
        检查 token 是否在黑名单中

        :param payload: JWT payload
        :return: 在黑名单中返回 True，否则返回 False
        """
        if not self.redis:
            return False
        jti = payload.get("jti")
        userid = payload.get("sub")
        if not jti:
            return False
        key = f"{self.blacklist_prefix}:{userid}:{jti}"
        exists = await self.redis.exists(key)
        return bool(exists)

    async def decode_token(
        self, token: str, expected_type: Optional[str] = None, options: Dict[str, bool] = None
    ) -> Optional[Dict[str, Any]]:
        """
        解码 JWT token

        :param token: JWT token
        :param expected_type: 期望的 token 类型，如果指定则会验证
        :return: 解码后的 payload，如果解码失败或类型不匹配则返回 None
        """
        options = options or {}
        payload = jwt.decode(token, self.secret_key, algorithms=[self.algorithm], options=options)

        if expected_type and payload.get("type") != expected_type:
            return None

        if await self._is_token_blacklisted(payload):
            # 可以添加日志记录
            logger.warning(f"Token is blacklisted: {token[:16]}...")
            return None

        return payload

    async def decode_access_token(self, token: str) -> Optional[Dict[str, Any]]:
        """
        解码访问 token

        :param token: 访问 token
        :return: 解码后的 payload，如果解码失败则返回 None
        """
        return await self.decode_token(token, "access")

    async def decode_refresh_token(self, token: str) -> Optional[Dict[str, Any]]:
        """
        解码刷新 token

        :param token: 刷新 token
        :return: 解码后的 payload，如果解码失败则返回 None
        """
        return await self.decode_token(token, "refresh")

    async def get_current_user_from_token(self, token: str) -> Optional[str]:
        """
        从访问 token 中获取当前用户标识

        :param token: 访问 token
        :return: 用户标识，如果解码失败则返回 None
        """
        payload = await self.decode_access_token(token)
        return payload.get("sub") if payload else None

    async def revoke_token(self, token: str) -> bool:
        """
        将 token 加入 Redis 黑名单，TTL = 剩余有效期（最多到原 exp）

        :param token: 要撤销的 token
        :return: 撤销成功返回 True，否则返回 False
        """
        if not self.redis:
            return False
        payload = jwt.decode(token, self.secret_key, algorithms=[self.algorithm])

        jti = payload.get("jti")
        exp = payload.get("exp")
        userid = payload.get("sub")
        if not jti or not exp:
            return False

        now = int(time.time())
        ttl = int(exp) - now
        if ttl <= 0:
            return False  # 已过期，无需加入黑名单

        key = f"{self.blacklist_prefix}:{userid}:{jti}"
        await self.redis.setex(key, ttl, "1")  # 值不重要，只关心 key 存在
        return True

    async def unrevoke_token(self, token: str) -> bool:
        """
        从 Redis 黑名单中移除 token（取消撤销）

        :param token: 要取消撤销的 token
        :return: 取消成功返回 True，否则返回 False
        """
        if not self.redis:
            return False
        payload = jwt.decode(token, self.secret_key, algorithms=[self.algorithm])

        jti = payload.get("jti")
        userid = payload.get("sub")
        if not jti:
            return False

        key = f"{self.blacklist_prefix}:{userid}:{jti}"
        result = await self.redis.delete(key)
        return result > 0  # 如果删除了至少一个键，则返回 True

    async def refresh_access_token(
        self, refresh_token: str, rotate_refresh: bool = False
    ) -> Optional[Tuple[str, Optional[str]]]:
        """
        使用刷新 token 刷新访问 token

        :param refresh_token: 刷新 token
        :param rotate_refresh: 是否轮换刷新 token
        :return: (新的访问 token, 新的刷新 token（如果 rotate_refresh 为 True）) 的元组，
                 如果刷新失败则返回 None
        """
        payload = await self.decode_refresh_token(refresh_token)
        if not payload:
            return None

        # 过滤掉 JWT 标准字段
        data = {k: v for k, v in payload.items() if k not in self.JWT_RESERVED_FIELDS}

        new_access = self.create_access_token(data)

        new_refresh = None
        if rotate_refresh:
            new_refresh = self.create_refresh_token(data)
            # 撤销旧 refresh token（可选，增强安全）
            await self.revoke_token(refresh_token)

        return new_access, new_refresh
